# -*- coding: utf-8 -*-
"""
# @file name    : L2_regularization.py
# @author       : QuZhang
# @date         : 2020-12-22 11:23
# @brief        : L2权值衰减使用试验
"""
import torch
import torch.nn as nn
from tools.common_tools import set_seed
from torch.utils.tensorboard import SummaryWriter
from matplotlib import pylab as plt
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

set_seed(1)
n_hidden = 200
max_iter = 2000
disp_interval = 200
lr_init = 0.01
writer = SummaryWriter(comment='_test_tensorboard', filename_suffix='12345678')


# ----------- step 1 数据 ----------
def gen_data(num_data=10, x_range=(-1, 1)):
    w = 1.5
    # *x_range：将元组拆成一个个独立的数据
    train_x = torch.linspace(*x_range, num_data).unsqueeze_(dim=1)
    # y = k * x + b_ , b为不同的值/噪声
    train_y = w*train_x + torch.normal(0, 0.5, size=train_x.size())
    test_x = torch.linspace(*x_range, num_data).unsqueeze_(1)
    test_y = w*test_x + torch.normal(0, 0.3, size=test_x.size())

    return train_x, train_y, test_x, test_y


# ----------- step 2 模型 ---------------
class MLP(nn.Module):
    def __init__(self, neural_num):
        super().__init__()
        self.linears = nn.Sequential(
            nn.Linear(1, neural_num),
            nn.ReLU(inplace=True),
            nn.Linear(neural_num, neural_num),
            nn.ReLU(inplace=True),
            nn.Linear(neural_num, neural_num),
            nn.ReLU(inplace=True),
            nn.Linear(neural_num, 1)
        )

    def forward(self, x):
        return self.linears(x)


if __name__ == '__main__':
    train_x, train_y, test_x, test_y = gen_data(num_data=10, x_range=(-1, 1))
    net_normal = MLP(neural_num=n_hidden)
    net_weight_decay = MLP(neural_num=n_hidden)

    # ------------- step 3 优化器 -----------------
    optim_normal = torch.optim.SGD(net_normal.parameters(), lr=lr_init, momentum=0.9)
    optim_wdecay = torch.optim.SGD(net_weight_decay.parameters(), lr=lr_init, momentum=0.9, weight_decay=0.01)

    # ------------- step 4 损失函数 -------------
    loss_func = torch.nn.MSELoss()

    # ------------ step 5 迭代训练 -------------
    for epoch in range(max_iter):

        # forward
        pred_normal, pred_wdecay = net_normal(train_x), net_weight_decay(train_x)
        # print(pred_normal.size())
        loss_normal, loss_wdecay = loss_func(pred_normal, train_y), loss_func(pred_wdecay, train_y)

        optim_normal.zero_grad()
        optim_wdecay.zero_grad()

        loss_normal.backward()
        loss_wdecay.backward()

        optim_normal.step()
        optim_wdecay.step()

        if (epoch+1) % disp_interval == 0:
            # 可视化
            # 遍历每层的参数
            for name, layer in net_normal.named_parameters():
                # name: 存储参数名(偏置，权重)
                # layer：存储参数对应的张量
                writer.add_histogram(name + '_grad_normal', layer.grad, epoch)  # 参数的梯度
                writer.add_histogram(name + '_data_normal', layer, epoch)  # 参数

            for name, layer in net_weight_decay.named_parameters():
                writer.add_histogram(name + '_grad_weight_decay', layer.grad, epoch)
                writer.add_histogram(name + "_data_weight_decay", layer, epoch)

            test_pred_normal, test_pred_wdecay = net_normal(test_x), net_weight_decay(test_x)

            # 绘图：观察预测test_pred_normal, test_pred_wdecay与真实值test_y的差距，判断是否过拟合
            # 散点图作为真实数据
            plt.scatter(train_x.data.numpy(), train_y.data.numpy(), c='blue', s=50, alpha=.3, label="train")
            plt.scatter(test_x.data.numpy(), test_y.data.numpy(), label="test", c='red', s=50, alpha=0.3)
            # 折线图作为回归数据
            plt.plot(test_x.data.numpy(), test_pred_normal.data.numpy(), 'r-', lw=3, label='no weight decay')
            plt.plot(test_x.data.numpy(), test_pred_wdecay.data.numpy(), 'b--', lw=3, label='weight decay')

            # 训练损失
            plt.text(-0.25, -1.5, 'no weight decay loss={:.6f}'.format(loss_normal.item()),
                     fontdict={'size': 15, 'color': 'red'})
            plt.text(-0.25, -2, 'weight decay loss={:.6f}'.format(loss_wdecay.item()),
                     fontdict={'size': 15, 'color': 'red'})

            plt.ylim((-2.5, 2.5))
            plt.legend(loc='upper left')
            plt.title("Epoch: {}".format(epoch + 1))
            plt.show()
            plt.close()
